"""
核心监控逻辑模块
使用psutil监控进程内存占用
"""
import psutil
import time
import sys
from typing import Optional, List, Tuple
from datetime import datetime

# 注意：Windows任务管理器显示的"内存"列默认是"专用工作集"（Private Working Set），
# 而psutil的rss返回的是"工作集"（Working Set），两者可能不同。
# 工作集 = 专用工作集 + 共享工作集
# 因此监控工具显示的值可能大于任务管理器显示的值。
# 如果需要与任务管理器完全一致，需要使用性能计数器API，实现较复杂。


class ProcessMonitor:
    """进程监控类"""
    
    def __init__(self, process_name: Optional[str] = None, process_id: Optional[int] = None):
        """
        初始化进程监控器
        
        Args:
            process_name: 进程名称（可选）
            process_id: 进程ID（可选）
        
        Raises:
            ValueError: 如果既没有提供进程名也没有提供进程ID
            ProcessNotFoundError: 如果找不到指定的进程
        """
        if not process_name and process_id is None:
            raise ValueError("必须提供进程名称或进程ID")
        
        self.process_name = process_name
        self.process_id = process_id
        self.processes = []
        self._find_processes()
    
    def _find_processes(self):
        """查找匹配的进程"""
        self.processes = []
        
        if self.process_id is not None:
            # 通过进程ID查找
            try:
                proc = psutil.Process(self.process_id)
                self.processes.append(proc)
            except psutil.NoSuchProcess:
                raise ProcessNotFoundError(f"找不到进程ID为 {self.process_id} 的进程")
        else:
            # 通过进程名查找
            found = False
            for proc in psutil.process_iter(['pid', 'name']):
                try:
                    if self.process_name.lower() in proc.info['name'].lower():
                        self.processes.append(psutil.Process(proc.info['pid']))
                        found = True
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    continue
            
            if not found:
                raise ProcessNotFoundError(f"找不到进程名为 '{self.process_name}' 的进程")
    
    def get_memory_info(self) -> List[Tuple[int, Optional[str], float, float]]:
        """
        获取当前内存信息
        
        注意：在Windows上，psutil的rss返回的是"工作集"（Working Set），
        而任务管理器默认显示的是"专用工作集"（Private Working Set）。
        工作集 = 专用工作集 + 共享工作集，因此监控工具显示的值可能大于任务管理器。
        
        Returns:
            进程信息列表，每个元素为 (process_id, process_name, memory_mb, memory_percent)
        """
        memory_info_list = []
        
        for proc in self.processes:
            try:
                memory_info = proc.memory_info()
                # 在Windows上，rss是工作集（Working Set）
                # 在Linux/Mac上，rss是常驻内存集（Resident Set Size）
                memory_mb = memory_info.rss / (1024 * 1024)  # 转换为MB
                memory_percent = proc.memory_percent()
                process_name = proc.name()
                
                memory_info_list.append((
                    proc.pid,
                    process_name,
                    memory_mb,
                    memory_percent
                ))
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                # 进程可能已经结束或没有权限访问
                continue
        
        return memory_info_list
    
    def monitor(self, interval: float = 1.0, duration: Optional[float] = None, 
                callback=None) -> List[Tuple[str, int, Optional[str], float, float]]:
        """
        持续监控进程内存
        
        Args:
            interval: 监控间隔（秒）
            duration: 监控持续时间（秒），如果为None则持续监控直到进程结束或手动停止
            callback: 回调函数，每次获取数据时调用 callback(timestamp, process_id, process_name, memory_mb, memory_percent)
        
        Returns:
            监控数据列表，每个元素为 (timestamp, process_id, process_name, memory_mb, memory_percent)
        """
        records = []
        start_time = time.time()
        
        try:
            while True:
                # 检查是否超过持续时间
                if duration and (time.time() - start_time) >= duration:
                    break
                
                # 重新查找进程（处理进程重启的情况）
                try:
                    self._find_processes()
                except ProcessNotFoundError:
                    # 进程已结束
                    break
                
                # 获取内存信息
                memory_info_list = self.get_memory_info()
                
                if not memory_info_list:
                    break
                
                timestamp = datetime.now().isoformat()
                
                for process_id, process_name, memory_mb, memory_percent in memory_info_list:
                    record = (timestamp, process_id, process_name, memory_mb, memory_percent)
                    records.append(record)
                    
                    # 调用回调函数
                    if callback:
                        callback(*record)
                
                time.sleep(interval)
        
        except KeyboardInterrupt:
            # 用户中断
            pass
        
        return records
    
    def is_running(self) -> bool:
        """检查进程是否还在运行"""
        try:
            self._find_processes()
            return len(self.processes) > 0
        except ProcessNotFoundError:
            return False


class ProcessNotFoundError(Exception):
    """进程未找到异常"""
    pass

